#!/usr/bin/env python
import sys, getopt, argparse, glob, math, os
import numpy as np
#from termcolor import colored

class bcolors:
    HEADER = '\033[95m'
    BLUE = '\033[94m'
    GREEN = '\033[92m'
    RED = '\033[31m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'
#	BLACK='\[\033[30m\]'
#	GREEN="\[\033[32m\]"
#	YELLOW="\[\033[33m\]"
    BLUE='\033[34m'
    MAGENTA='\033[35m'
    CYAN='\033[36m'
    WHITE='\033[37m'

source_dir="."
out_dir="./merged/"

verbose= False

merge_mode="by bin index"
merge_mode="by bin X value"

printouthisto="none"
#printouthisto='"m(W+_candidate)lin"'
#printouthisto='"m(5)"'
#printouthisto='"sqrt(p^2(5))"'
#printouthisto='"E(b_and_bbar)30_150_4"'
#printouthisto='"min(mlb)"'
#printouthisto='"sttj_240_340"'

print 'printing out the merging of the histogram ', bcolors.UNDERLINE, printouthisto, bcolors.ENDC

parser = argparse.ArgumentParser(description='Process some gnuplot scripts to merge histograms.')
parser.add_argument('jobID', metavar='N', nargs='+',  help='an integer for the jobID to be treated')
args = parser.parse_args()

#jobID="2523364"
jobID=args.jobID
print "Merging histograms from job:", bcolors.MAGENTA, str(jobID[0]),  bcolors.ENDC

file_list = glob.glob(source_dir + '/' + '*' + jobID[0] + '*.gnu')

base_file_name=str(file_list[0]).split(str(jobID[0]))
#  + jobID[0] +
out_filename= str(base_file_name[0])  + str(jobID[0]) + "_x" + str(len(file_list)) + ".gnu"
print bcolors.RED, ' will output to file ', bcolors.ENDC, out_filename

histogramsfromallfiles = {}

for file_path in file_list:

	f = open(file_path,'r')
	#f = open("tt_bbl_lord_mstw8nl_85__85__LHC14_2.0M_mt170_pTj30_r4_eta2.4_cut1.gnu",'r')


	titles=[]
	duplicate_titles=[]
	DataIsOpen = False
	thishistogram=[]
	histogramsfromthisfile = {}

	counter = 0
	for line in f:
		counter+=1
#		if counter < 200:
		if True:
	#	print 'line ',counter,':', line
			uncommentedline	= line.split('#')
	#	print 'uncommented', uncommentedline[0]
	#	fields	= ' '.join(line.split())
			if len(uncommentedline[0]) > 0:
				fields	= uncommentedline[0].split()
	#		print 'fields', fields
				if len(fields)>0:

					if fields[0] == "e":
						DataIsOpen = False
						#print "closing the histogram", titles[-1]
						title=titles[-1]
						histogramsfromthisfile[title]=thishistogram
						if titles[-1] == printouthisto:
					    		print "histogram:", histogramsfromthisfile[printouthisto]
						thishistogram=[]
	#				print fields[0]
	#				print "printing the histogram", titles[-1]		
	#				print thishistogram
	#				dictionarylemma={ title: thishistogram}
					if DataIsOpen == True:
	#write x y dy into an array that knows about titles[-1]
						thispoint=[fields[0], fields[1],fields[2]]
						thishistogram.append(thispoint)
						if titles[-1] == printouthisto:
					    		print "added ", thispoint, "to ", titles[-1], " from ",file_path 
					if fields[0] == "set" and fields[1]=="xlabel":
						#print "found histogram", fields[2]
						if fields[2] in titles:
							duplicate_titles.append(fields[2])
							if verbose:	
								print 'so far the following titles have been used multiple times: ', duplicate_titles
						titles.append(fields[2])
					if fields[0] == "plot":
						DataIsOpen = True
						#print "beginning histogram ", titles[-1]
						thishistogram=[]

#	print histogramsfromthisfile				 	
#	print 'Found histogram names ', titles
#	print 'Stored histogram names ', histogramfromthisfile.keys()
	histogramsfromallfiles[file_path]=histogramsfromthisfile
#print histogramsfromallfiles[:]

mergedhistograms = {}



for name in titles:
	if name == printouthisto:
#	if True:	
 		print 'merging ', name, ' of ', len(histogramsfromallfiles[file_list[0]][name]), ' bins'
 	thismergedhisto=[]
	for bin_index in range(0,len(histogramsfromallfiles[file_list[0]][name])):
		# we are merging the bin bin_index
		total_weight=0.0
		pre_weighted_average=0.0
		if name == printouthisto:
			print 'bin#:',bin_index
		for file_index in range (0,len(file_list)):
			# fixed the bin we are looking at, now I look for it in all the files
 			#print file_index
			number_of_bins_this_file=len(histogramsfromallfiles[file_list[file_index]][name])
			if name == printouthisto:
				print 'file#:', file_index+1,'/',len(file_list), ' ' ,file_list[file_index],' of ',number_of_bins_this_file, ' bins'
			x_orig=float(histogramsfromallfiles[file_list[0]][name][bin_index][0])
			#
			# start the procedure to get the x,y,dy from the x_orig
			#
			if merge_mode=="by bin index":
				if bin_index < number_of_bins_this_file:
					#only files where the bin is filled are included

					x=float(histogramsfromallfiles[file_list[file_index]][name][bin_index][0])
					# print x, "=" ,x_orig,"?"
					y=float(histogramsfromallfiles[file_list[file_index]][name][bin_index][1])
					dy=float(histogramsfromallfiles[file_list[file_index]][name][bin_index][2])
					if x_orig != x:
						if verbse:
							 print 'in file #',file_index,":", file_list[file_index]
							 print name,"[bin=",bin_index,"]=",x, "!=" ,x_orig,"?"
				# at this point the x,y,dy are valorized and the checks and the combination can take place
			if merge_mode=="by bin X value":
				for bin_ith in range (0,number_of_bins_this_file):
					thisX=float(histogramsfromallfiles[file_list[file_index]][name][bin_ith][0])
					if thisX == x_orig:
						x=float(histogramsfromallfiles[file_list[file_index]][name][bin_ith][0])
						y=float(histogramsfromallfiles[file_list[file_index]][name][bin_ith][1])
						dy=float(histogramsfromallfiles[file_list[file_index]][name][bin_ith][2])

			# there is the chance that the point has not been found
			if 'x' in locals():
				# check that x is the same as the ones before!
				if x_orig == x:
					# if one histogram is misaligned (it starts at a different bin) its data is discarded
					# if one histogram end before the other only the data till its end is used in combination
					if name == printouthisto:
						print "merging in the point ", x," ",y," ",dy, "originally from", file_list[file_index]
					pre_weighted_average+=(y)/math.pow(dy,2.0)
					total_weight+=1.0/math.pow(dy,2.0)
				del x 
				del y
				del dy	
			#end loop on files
		sigma_combined=math.sqrt(1.0/total_weight)
		weighted_average=pre_weighted_average/total_weight
		x_bin=float(histogramsfromallfiles[file_list[0]][name][bin_index][0])
		thispoint=[x_bin, weighted_average, sigma_combined]
		if name == printouthisto:
			print bcolors.BLUE, 'after combination:',  bcolors.ENDC, 	thispoint
		thismergedhisto.append(thispoint)
	    #done with loop on bins
	mergedhistograms[name]=thismergedhisto
	if name == printouthisto:
		print mergedhistograms[printouthisto]		
	#done loop on observables
#	print 'merged ', name, 'from', len(file_list) , 'files'



#for arg in sys.argv:
#    print arg
print bcolors.BLUE, 'files analyzed: ',  bcolors.ENDC , histogramsfromallfiles.keys()
print bcolors.BLUE, 'histograms read:',  bcolors.ENDC, histogramsfromallfiles[file_list[0]].keys()
print bcolors.BLUE, 'histograms merged: ',  bcolors.ENDC, mergedhistograms.keys()


print bcolors.RED, ' will output to file ', bcolors.ENDC, out_filename

f = open(file_list[0],'r')
print "Opening the file ", out_filename, " in folder ", out_dir
print(os.path.isdir(out_dir))
overwritten_output=False
if os.path.exists(out_dir+out_filename):
	print bcolors.RED, "output file is overwritten", bcolors.ENDC
	overwritten_output=True
	
target = open(out_dir+out_filename, 'w')
DataIsOpen = False

written_titles = []

#ignore_titles = [ '"Et(5)*cosh(y)"' ]
ignore_titles = []

WriteData=False
Skip=False

counter = 0
for line in f:
	counter+=1
#		if counter < 200:
	if True:
		uncommentedline	= line.split('#')
		if DataIsOpen == False: 
			target.write(line)
		if len(uncommentedline[0]) > 0:
			fields	= uncommentedline[0].split()
			if len(fields)>0:

				if fields[0] == "e":
					DataIsOpen = False
					Skip=False
					target.write(line)
				if DataIsOpen == True and WriteData == True and Skip == False:
					# write the histogram from the memory
					if verbose: 
						print 'writing to file the merged histogram ', written_titles[-1]
					for item in mergedhistograms[written_titles[-1]]:
						#print item
						#target.write("%s\n" % item[:])
						text="     "
						for elem in item: 
							text += str(elem) + "        " # add commas if you wish 
						target.write( text + "\n" )
					# once written, close the stream
					WriteData=False
				if fields[0] == "set" and fields[1]=="xlabel":
					if fields[2] in ignore_titles:
						print "skipping histogram", fields[2]
						Skip=True
					else:
						written_titles.append(fields[2])
				if fields[0] == "plot":
					DataIsOpen = True
					WriteData = True
					#print "beginning histogram ", titles[-1]
					
print bcolors.BLUE, 'merged histograms written: ',  bcolors.ENDC, written_titles
print 'the following titles have been used multiple times: ', duplicate_titles

if overwritten_output:
	print bcolors.RED, "output file existed and was overwritten", bcolors.ENDC
